import numpy as np
import matplotlib.pyplot as plt
from itertools import product, combinations
times = np.loadtxt('time3.5.txt')

idx = 1
t = np.linspace(0,500,50001)
a = np.loadtxt('Li%s.txt'%(str(idx)))

fig = plt.figure(figsize=(8,6))
ax = plt.axes(projection='3d')
ax.set_box_aspect([1,1,1])
#plt.subplots_adjust(left=0, right=0.8, bottom=0, top=0.8)

ax.plot3D(a[:,0],a[:,1],a[:,2],lw=0.1,zorder=0)
ax.scatter(a[0,0],a[0,1],a[0,2],color='blue',label='_nolegend_')

#r1 = ((t[:] >= 104.5) & (t[:] <= 105))
#r2 = ((t[:] >= 113.25) & (t[:] <= 113.75))
#r3 = ((t[:] >= 144.5) & (t[:] <= 145))
#r4 = ((t[:] >= 207.5) & (t[:] <= 208))
#r5 = ((t[:] >= 210) & (t[:] <= 210.5))

ax.plot3D(a[:,0],a[:,1],a[:,2],lw=1,color='green',zorder=1)
for i in range(len(times)):
#    r = globals()['r%s' % i]
    r = ((t[:] >= times[i]-0.1) & (t[:] <= times[i]+0.1))
    ax.plot3D(a[r,0],a[r,1],a[r,2],lw=2,color='red',zorder=1)

a = 2.0104010492000000e+01/2
# draw cube
r = [0, a]
for s, e in combinations(np.array(list(product(r, r, r))), 2):
    if np.sum(np.abs(s-e)) == r[1]-r[0]:
        ax.plot3D(*zip(s, e), color="black",lw=0.5)

ax.set_axis_off()
#ax.text(5,1,19,r'0$\rightarrow$144.5 ps',fontsize=16,color='green',ha='left')
#ax.text(5,1,17.5,r'144.5$\rightarrow$145 ps',fontsize=16,color='orange',ha='left')
#ax.text(5,1,16,r'145$\rightarrow$397 ps',fontsize=16,color='limegreen',ha='left')
#ax.text(5,1,14.5,r'397$\rightarrow$397.5 ps',fontsize=16,color='red',ha='left')
#ax.text(5,1,13,r'397.5$\rightarrow$500 ps',fontsize=16,color='forestgreen',ha='left')
ax.xaxis.set_rotate_label(False) 
ax.yaxis.set_rotate_label(False)
ax.zaxis.set_rotate_label(False)
ax.set_xlabel(r'x ($\rm{\AA}$)',fontsize=16,labelpad=6,rotation=0)
ax.set_ylabel(r'y ($\rm{\AA}$)',fontsize=16,labelpad=6,rotation=0)
ax.set_zlabel(r'z ($\rm{\AA}$)',fontsize=16,labelpad=6,rotation=90)
ax.grid(False)
ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
ax.tick_params(direction='in',labelsize=16,pad=2)
#ax.view_init(8, -60)
ax.view_init(9, -45)
plt.savefig('TrajLi%s_range_Dec23.pdf'%(str(idx)),format='pdf')
plt.savefig('TrajLi%s_range_Dec23.eps'%(str(idx)),format='eps')
plt.savefig('TrajLi%s_range_Dec23.png'%(str(idx)),format='png',dpi=900)
plt.show()
